{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "b7d716a8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2.9.0+cpu'"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "torch.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "793cb1cf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2.3.3'"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "np.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "29800cd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "3f361093",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[5.1487e-01, 8.9170e-01, 7.3196e-01, 7.0828e-01],\n",
      "        [2.0453e-01, 6.2025e-01, 8.5274e-01, 1.1414e-04],\n",
      "        [5.6524e-01, 2.9085e-01, 5.7118e-01, 7.7456e-01]])\n"
     ]
    }
   ],
   "source": [
    "random_torch=torch.rand(3,4)\n",
    "print(random_torch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "364c22bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TokenEmbedding(nn.Embedding):\n",
    "    def __init__(self, vocab_size, d_model):\n",
    "        super(TokenEmbedding, self).__init__(vocab_size, d_model, padding_idx=1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "4f7a70b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionalEmbedding(nn.Module):\n",
    "    def __init__(self, d_model, max_len=5000, device='cpu'):\n",
    "        super(PositionalEmbedding, self).__init__()\n",
    "        self.encoding = torch.zeros(max_len, d_model).to(device)\n",
    "        self.encoding.requires_grad = False\n",
    "        pos = torch.arange(0, max_len).to(device)\n",
    "        pos=pos.float().unsqueeze(dim=1)\n",
    "        _zi = torch.arange(0, d_model, step=2).float().to(device)\n",
    "        self.encoding[:, 0::2] = torch.sin(pos / (10000 ** (_zi / d_model)))\n",
    "        self.encoding[:, 1::2] = torch.cos(pos / (10000 ** (_zi / d_model)))\n",
    "\n",
    "    def forward(self, x):\n",
    "        batch_size, seq_len = x.size()\n",
    "        return self.encoding[:seq_len, :]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "f113803c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerEmbedding(nn.Module):\n",
    "    def __init__(self, vocab_size, d_model, max_len, drop_prob, device='cpu'):\n",
    "        super(TransformerEmbedding, self).__init__()\n",
    "        self.tok_emb = TokenEmbedding(vocab_size, d_model)\n",
    "        self.pos_emb = PositionalEmbedding(d_model, max_len=max_len, device=device)\n",
    "        self.dropout = nn.Dropout(p=drop_prob)\n",
    "\n",
    "    def forward(self, x):\n",
    "        tok_emb = self.tok_emb(x)\n",
    "        pos_emb = self.pos_emb(x)\n",
    "        return self.dropout(tok_emb + pos_emb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64441341",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "45852272",
   "metadata": {},
   "outputs": [],
   "source": [
    "q = torch.rand(128,32,512)\n",
    "k = torch.rand(128,32,512)\n",
    "v = torch.rand(128,32,512)\n",
    "d_model=512\n",
    "n_head=8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47c9cfa7",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, d_model, n_head):\n",
    "        super(MultiHeadAttention, self).__init__()\n",
    "        self.d_model = d_model\n",
    "        self.n_head = n_head\n",
    "        self.w_q = nn.Linear(d_model, n_head)\n",
    "        self.w_k = nn.Linear(d_model, n_head)\n",
    "        self.w_v = nn.Linear(d_model, n_head)\n",
    "        self.w_combine = nn.Linear(d_model, n_head)\n",
    "        self.softmax = nn.Softmax(dim=-1)\n",
    "\n",
    "    def forward(self, q, k, v, mask=None):\n",
    "        batch, time, demension=q.size()\n",
    "        n_d = self.d_model//self.n_head\n",
    "\n",
    "        # Linear projections\n",
    "        q = self.w_q(q).view(batch, time, self.n_head, n_d).permute(0, 2, 1, 3)\n",
    "        k = self.w_q(k).view(batch, time, self.n_head, n_d).permute(0, 2, 1, 3)\n",
    "        v = self.w_q(v).view(batch, time, self.n_head, n_d).permute(0, 2, 1, 3)\n",
    "\n",
    "        score = q@k.transpose(2,3)/math.sqrt(n_d)\n",
    "        if mask is not None:\n",
    "            score = score.masked_fill(mask == 0, -1e9)\n",
    "        score=self.softmax(score)@v\n",
    "        score = score.permute(0, 2, 1, 3).contiguous().view(batch, time, demension)\n",
    "        \n",
    "\n",
    "        output = self.w_combine(score)\n",
    "\n",
    "        return output\n",
    "\n",
    "attention: MultiHeadAttention = MultiHeadAttention(d_model=512, n_head=8)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "2f2b8e19",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([128, 32, 512]), torch.Size([128, 32, 8]))"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xx = attention.w_q(q)\n",
    "q.shape, xx.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "282e3384",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "6af91d3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "class LayerNorm(nn.Module):\n",
    "    def __init__(self, d_model, aps=1e-12):\n",
    "        super().__init__()\n",
    "        self.gama=nn.Parameter(torch.ones(d_model))\n",
    "        self.beta=nn.Parameter(torch.zeros(d_model))\n",
    "        self.aps = aps\n",
    "    \n",
    "    def forward(self, x):\n",
    "        mean = x.mean(dim=-1, keepdim=True)\n",
    "        # std = x.std(dim=-1, keepdim=True)\n",
    "        # norm_x = (x - mean) / (std + self.aps)\n",
    "        # return self.gama * norm_x + self.beta\n",
    "        var = x.var(dim=-1, keepdim=True, unbiased=False)\n",
    "        norm_x = (x - mean) / torch.sqrt(var + self.aps)\n",
    "        return self.gama * norm_x + self.beta\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "368a5bac",
   "metadata": {},
   "source": [
    "-----------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "0f716f15",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionWiseFeedForward(nn.Module):\n",
    "    def __init__(self, d_model, hidden, dropout=0.1):\n",
    "        super().__init__()\n",
    "        self.fc1 = nn.Linear(d_model, hidden)\n",
    "        self.fc2 = nn.Linear(hidden, d_model)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.fc1(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.dropout(x)\n",
    "        x= self.fc2(x)\n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "457c441e",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class EncoderLayer(nn.Module):\n",
    "    def __init__(self, d_model, ffn_hidden, n_head, dropout=0.1):\n",
    "        super(EncoderLayer, self).__init__()\n",
    "        self.self_attn = MultiHeadAttention(d_model, n_head)\n",
    "        self.normal= LayerNorm(d_model)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.ffn = PositionWiseFeedForward(d_model, ffn_hidden, dropout=dropout)\n",
    "        self.normal2 = LayerNorm(d_model)\n",
    "        self.dropout2 = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x, mask):\n",
    "        _x=x\n",
    "        x = self.self_attn(x, x, x, mask)\n",
    "        x = self.dropout(x)\n",
    "        x = self.normal(x+_x) # 残差连接\n",
    "        _x = x\n",
    "        x = self.ffn(x)\n",
    "        x = self.dropout2(x)\n",
    "        x=self.normal2(x+_x)\n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "0d1bced8",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    def __init__(self, enc_vocab_size, max_len, d_model, ffn_hidden, n_head, n_layer, dropout=0.1):\n",
    "        super(Encoder, self).__init__()\n",
    "        self.embedding = TransformerEmbedding(enc_vocab_size, max_len, d_model, dropout )\n",
    "        self.layers = nn.ModuleList([\n",
    "            EncoderLayer(d_model, ffn_hidden, n_head, dropout)\n",
    "            for _ in range(n_layer)\n",
    "        ])\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.embedding(x)\n",
    "        for layer in self.layers:\n",
    "            x = layer(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "694b0dde",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "c7eb27e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydoc import doc\n",
    "\n",
    "\n",
    "class DecoderLayer(nn.Module):\n",
    "    def __init__(self, d_model, ffn_hidden, n_head,dropout_prob=0.1):\n",
    "        super(DecoderLayer, self).__init__()\n",
    "        self.attention1=MultiHeadAttention(d_model, n_head)\n",
    "        self.norm1=LayerNorm(d_model)\n",
    "        self.dropout1=nn.Dropout(dropout_prob)\n",
    "        # 跨模态注意力机制\n",
    "        self.cross_attention=MultiHeadAttention(d_model, n_head)\n",
    "        self.norm2=LayerNorm(d_model)\n",
    "        self.dropout2=nn.Dropout(dropout_prob)\n",
    "        self.ffn=PositionWiseFeedForward(d_model, ffn_hidden, dropout_prob)\n",
    "        self.norm3=LayerNorm(d_model)\n",
    "        self.dropout3=nn.Dropout(dropout_prob)\n",
    "\n",
    "    def forward(self, dec, enc, t_mask,s_mask):\n",
    "        _x= dec\n",
    "        x=self.attention1(dec, dec, dec, t_mask)\n",
    "        x=self.dropout1\n",
    "        x = self.norm1(x+_x)\n",
    "        _x = x\n",
    "        x = self.cross_attention(x, enc, enc, s_mask)\n",
    "        x = self.dropout2(x)\n",
    "        x = self.norm2(x+_x)\n",
    "        x=self.ffn(x)\n",
    "        x=self.dropout3(x)\n",
    "        x=self.norm3(x+_x)\n",
    "        return x\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "475febe9",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    def __init__(self, dec_voc_size, max_len, d_model, ffn_hidden, n_heads, n_layers, dropout)  :   \n",
    "        super().__init__()\n",
    "        self.embedding = TransformerEmbedding(dec_voc_size, max_len, d_model, dropout)\n",
    "        self.layers = nn.ModuleList([DecoderLayer(d_model, ffn_hidden, n_heads, dropout) for _ in range(n_layers)])\n",
    "        self.fc = nn.Linear(d_model, dec_voc_size)\n",
    "    \n",
    "    def forward(self, dec, enc, t_mask, s_mask):\n",
    "        dec = self.embedding(dec)\n",
    "        for layer in self.layers:\n",
    "            dec = layer(dec, enc, t_mask, s_mask)\n",
    "\n",
    "        dec = self.fc(dec)\n",
    "        return dec"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9cb4a3a6",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0f67ee7",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Transformer(nn.Module):\n",
    "    def __init__(self, src_pad_idx, trg_pad_idx, enc_voc_size, dec_voc_size, d_model,max_len, n_heads,ffn_hidden,n_layers, dropout_prob):\n",
    "        super(Transformer, self).__init__()\n",
    "        self.encoder = Encoder(enc_voc_size, max_len=max_len, d_model=d_model, n_head=n_heads, ffn_hidden=ffn_hidden, n_layer= n_layers, dropout= dropout_prob)\n",
    "        self.decoder = Decoder(dec_voc_size, max_len=max_len, d_model=d_model, n_heads=n_heads, ffn_hidden=ffn_hidden, n_layers= n_layers, dropout= dropout_prob)\n",
    "        self.src_pad_idx = src_pad_idx\n",
    "        self.trg_pad_idx = trg_pad_idx\n",
    "    \n",
    "    def mak_pad_mask(self,q,k,pad_idx_q, pad_idx_k):\n",
    "        len_q, len_k=q.size(1), k.size(1)\n",
    "        q=q.ne(pad_idx_q).unsqueeze(1).unsqueeze(3)\n",
    "        q=q.repeat(1,1,1,len_k)\n",
    "        k=k.ne(pad_idx_k).unsqueeze(1).unsqueeze(2)\n",
    "        k=k.repeat(1,1,len_q,1)\n",
    "        mask=q & k\n",
    "        return mask\n",
    "    \n",
    "    def make_casual_mask(self,q,k):\n",
    "        mask=torch.tril(torch.ones(q.size(1), k.size(1))).type(torch.BoolTensor)\n",
    "        return mask\n",
    "    \n",
    "    def forward(self, src, trg):\n",
    "        src_mask = self.mak_pad_mask(src, src, self.src_pad_idx, self.src_pad_idx)\n",
    "        trg_mask = self.mak_pad_mask(trg, trg, self.trg_pad_idx, self.trg_pad_idx) * self.make_casual_mask(trg, trg)\n",
    "        enc=self.encoder(src, src_mask)\n",
    "        dec=self.decoder(trg, enc, trg_mask, src_mask)\n",
    "        return dec\n",
    "\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".conda",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
